import random
import numpy as np
import matplotlib.pyplot as plt
from deap import base, creator, tools, algorithms


# ========== 1. 定义目标函数 ==========
def evaluate(individual):
    """双目标优化问题：
    f1 = x² + y²
    f2 = (x-1)² + (y-1)²"""
    x, y = individual
    f1 = x ** 2 + y ** 2  # 目标1：最小化到原点的距离
    f2 = (x - 1) ** 2 + (y - 1) ** 2  # 目标2：最小化到点(1,1)的距离
    return f1, f2  # 返回元组


# ========== 2. 创建类型和工具箱 ==========
creator.create("FitnessMulti", base.Fitness, weights=(-1.0, -1.0))  # 双目标最小化
creator.create("Individual", list, fitness=creator.FitnessMulti)

toolbox = base.Toolbox()
toolbox.register("attr_float", random.uniform, -10, 10)  # 变量范围[-10,10]
toolbox.register("individual", tools.initRepeat, creator.Individual,
                 toolbox.attr_float, n=2)  # 2维个体(x,y)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)

# ========== 3. 配置遗传算法操作 ==========
toolbox.register("evaluate", evaluate)
toolbox.register("mate", tools.cxBlend, alpha=0.5)  # 混合交叉
toolbox.register("mutate", tools.mutGaussian, mu=0, sigma=1, indpb=0.2)
toolbox.register("select", tools.selNSGA2)  # 多目标选择算子


# ========== 4. 运行算法 ==========
def main():
    pop = toolbox.population(n=200)  # 调整种群大小为200
    hof = tools.ParetoFront()  # 帕累托前沿存储器

    stats = tools.Statistics(lambda ind: ind.fitness.values)
    stats.register("avg", np.mean, axis=0)
    stats.register("std", np.std, axis=0)

    # 运行NSGA-II算法，调整交叉概率、变异概率和迭代次数
    pop, logbook = algorithms.eaMuPlusLambda(
        pop, toolbox, mu=200, lambda_=200,
        cxpb=0.8, mutpb=0.2, ngen=80,
        stats=stats, halloffame=hof, verbose=True
    )
    return pop, hof, logbook


if __name__ == "__main__":
    pop, hof, log = main()

    # ========== 5. 可视化结果 ==========
    # 绘制帕累托前沿
    front = np.array([ind.fitness.values for ind in hof])
    plt.scatter(front[:, 0], front[:, 1], c="red", s=30, label="Pareto Front")

    # 绘制初始种群和最终种群对比
    init_pop = np.array([ind for ind in toolbox.population(n=200)])
    final_pop = np.array([ind for ind in pop])

    plt.scatter(init_pop[:, 0] ** 2 + init_pop[:, 1] ** 2,
                (init_pop[:, 0] - 1) ** 2 + (init_pop[:, 1] - 1) ** 2,
                c="blue", alpha=0.3, label="Initial Population")
    plt.scatter(final_pop[:, 0] ** 2 + final_pop[:, 1] ** 2,
                (final_pop[:, 0] - 1) ** 2 + (final_pop[:, 1] - 1) ** 2,
                c="green", alpha=0.7, label="Final Population")

    plt.xlabel("f1(x,y) = x² + y²")
    plt.ylabel("f2(x,y) = (x-1)² + (y-1)²")
    plt.title("NSGA-II Multiobjective Optimization")
    plt.legend()
    plt.grid()
    plt.show()